#include "cld_reg.h"

#include <stdint.h>
#include <time.h>  // clock(), CLOCKS_PER_SEC

#include <iostream>
#include <string>
#include <iomanip>  // io manipulator, set cin or cout format, std::fixed, std::setprecision()

#include <pcl/filters/voxel_grid.h>  // pcl::VoxelGrid
#include <pcl/search/kdtree.h>  // pcl::search::KdTree
#include <pcl/keypoints/iss_3d.h>  // pcl::ISSKeypoint3D
#include <pcl/features/normal_3d.h>  // pcl::NormalEstimation
#include <pcl/features/fpfh.h>  // pcl::FPFHEstimation
#include <pcl/registration/gicp.h>  // pcl::GeneralizedIterativeClosestPoint

#include "ia_gror.h"

namespace cldreg {

Reg::Reg(pcl::PointCloud<PointType>::Ptr& cldptr_bg, double d_resolution, int n_corr_max) {
	d_resolution_ = d_resolution;
	n_corr_max_ = n_corr_max;
	cldptr_bg_ds_ = boost::make_shared<pcl::PointCloud<PointType> >();
	cldptr_bg_iss_ = boost::make_shared<pcl::PointCloud<PointType> >();
	indptr_bg_iss_ = boost::make_shared<pcl::PointIndices>();
	cldptr_bg_fpfh_ = boost::make_shared<pcl::PointCloud<pcl::FPFHSignature33> >();
	prepCld(cldptr_bg, d_resolution_, cldptr_bg_ds_, cldptr_bg_iss_, indptr_bg_iss_, cldptr_bg_fpfh_);
}

bool Reg::getFinalTransformation(pcl::PointCloud<PointType>::Ptr& cldptr_raw, Eigen::Matrix4f& mat_res) {
	mat_res = Eigen::Matrix4f::Identity();

	pcl::PointCloud<PointType>::Ptr cldptr_raw_ds(new pcl::PointCloud<PointType>);
	pcl::PointCloud<PointType>::Ptr cldptr_raw_iss(new pcl::PointCloud<PointType>);
	pcl::PointIndicesPtr indptr_raw_iss(new pcl::PointIndices);
	pcl::PointCloud<pcl::FPFHSignature33>::Ptr cldptr_raw_fpfh(new pcl::PointCloud<pcl::FPFHSignature33>());
	prepCld(cldptr_raw, d_resolution_, cldptr_raw_ds, cldptr_raw_iss, indptr_raw_iss, cldptr_raw_fpfh);

	pcl::CorrespondencesPtr corrptr(new pcl::Correspondences);
	std::vector<int> vec_corr_bg;
	std::vector<int> vec_corr_raw;
	correspondenceSearch(cldptr_raw_fpfh, cldptr_bg_fpfh_, corrptr, n_corr_max_, vec_corr_raw, vec_corr_bg);

	pcl::PointCloud<PointType>::Ptr cldptr_gror(new pcl::PointCloud<PointType>);
	Eigen::Matrix4f mat_gror = Eigen::Matrix4f::Identity();
	double d_score_gror;
	grorReg(cldptr_raw_iss, cldptr_bg_iss_, corrptr, d_resolution_, cldptr_gror, mat_gror, d_score_gror);
	//if (d_score_gror > 20.0) return false;

	pcl::PointCloud<PointType>::Ptr cldptr_gicp(new pcl::PointCloud<PointType>);
	Eigen::Matrix4f mat_gicp = Eigen::Matrix4f::Identity();
	double d_score_gicp;
	gicpReg(cldptr_raw_ds, cldptr_bg_ds_, mat_gror, cldptr_gicp, mat_gicp, d_score_gicp);
	//if (d_score_gicp > 2.0) return false;

	mat_res = mat_gicp;
	return true;
}

void Reg::prepCld(	pcl::PointCloud<PointType>::Ptr& cldptr_src, 
					double d_resolution, 
					pcl::PointCloud<PointType>::Ptr& cldptr_ds, 
					pcl::PointCloud<PointType>::Ptr& cldptr_iss, 
					pcl::PointIndicesPtr& indptr_iss, 
					pcl::PointCloud<pcl::FPFHSignature33>::Ptr& cldptr_fpfh) {
	clock_t t0, t1, t2;
	t0 = clock();

	size_t sz0 = cldptr_src->size();

	// voxel grid downsample, cldptr_ds
	pcl::VoxelGrid<PointType> vg;
	vg.setLeafSize(d_resolution, d_resolution, d_resolution);
	vg.setInputCloud(cldptr_src);
	vg.filter(*cldptr_ds);
	size_t sz1 = cldptr_ds->size();
	std::cout << std::fixed << std::setprecision(1) << "voxel grid downsample, " << d_resolution << ": " << sz0 << " -> " << sz1 << std::endl;

	// extract iss keypoints, cldptr_iss, indptr_iss
	double iss_salient_radius_ = d_resolution * 6;
	double iss_non_max_radius_ = d_resolution * 2;  // default 4, office 2, railway 9
	double iss_gamma_21_(0.975);
	double iss_gamma_32_(0.975);
	double iss_min_neighbors_(4);
	int iss_threads_(1);
	pcl::search::KdTree<PointType>::Ptr kdptr_iss(new pcl::search::KdTree<PointType>());
	pcl::ISSKeypoint3D<PointType, PointType> iss_detector;
	iss_detector.setSearchMethod(kdptr_iss);
	iss_detector.setSalientRadius(iss_salient_radius_);
	iss_detector.setNonMaxRadius(iss_non_max_radius_);
	iss_detector.setThreshold21(iss_gamma_21_);
	iss_detector.setThreshold32(iss_gamma_32_);
	iss_detector.setMinNeighbors(iss_min_neighbors_);
	iss_detector.setNumberOfThreads(iss_threads_);
	iss_detector.setInputCloud(cldptr_ds);
	iss_detector.compute(*cldptr_iss);
	indptr_iss->indices = iss_detector.getKeypointsIndices()->indices;
	indptr_iss->header = iss_detector.getKeypointsIndices()->header;

	// compute normals, cldptr_norm
	pcl::PointCloud<pcl::Normal>::Ptr cldptr_norm(new pcl::PointCloud<pcl::Normal>());
	pcl::search::KdTree<PointType>::Ptr kdptr_norm(new pcl::search::KdTree<PointType>());
	pcl::NormalEstimation<PointType, pcl::Normal> ne;
	ne.setInputCloud(cldptr_ds);
	ne.setSearchMethod(kdptr_norm);
	ne.setRadiusSearch(d_resolution * 3);
	ne.compute(*cldptr_norm);

	// compute fpfhs using normals, cldptr_fpfh
	pcl::search::KdTree<PointType>::Ptr kdptr_fpfh(new pcl::search::KdTree<PointType>());
	pcl::FPFHEstimation<PointType, pcl::Normal, pcl::FPFHSignature33> fpfh_est;
	fpfh_est.setInputCloud(cldptr_ds);
	fpfh_est.setInputNormals(cldptr_norm);
	fpfh_est.setSearchMethod(kdptr_fpfh);
	fpfh_est.setRadiusSearch(d_resolution * 8);
	fpfh_est.setIndices(indptr_iss);
	fpfh_est.compute(*cldptr_fpfh);

	t2 = clock();
	//std::cout << std::endl;
	std::cout << std::fixed << std::setprecision(3) << "----prepCld (s): " << double(t2 - t0) / CLOCKS_PER_SEC << std::endl;
}

void Reg::correspondenceSearch(	pcl::PointCloud<pcl::FPFHSignature33>::Ptr& cldptr_src_fpfh, 
								pcl::PointCloud<pcl::FPFHSignature33>::Ptr& cldptr_tgt_fpfh, 
								pcl::CorrespondencesPtr& corrptr, 
								const int n_corr_max, 
								std::vector<int>& vec_corr_src, 
								std::vector<int>& vec_corr_tgt) {
	clock_t t1, t2;
	t1 = clock();

	int n = std::min(n_corr_max, (int)cldptr_tgt_fpfh->size());
	corrptr->clear();
	vec_corr_src.assign(cldptr_src_fpfh->size(), 0);
	vec_corr_tgt.assign(cldptr_tgt_fpfh->size(), 0);
	pcl::KdTreeFLANN<pcl::FPFHSignature33> kd_src;
	kd_src.setInputCloud(cldptr_src_fpfh);
	pcl::KdTreeFLANN<pcl::FPFHSignature33> kd_tgt;
	kd_tgt.setInputCloud(cldptr_tgt_fpfh);
	for (size_t i = 0; i < cldptr_src_fpfh->size(); i++) {
		std::vector<int> vec_idx_tgt(n);
		std::vector<float> vec_dist_tgt(n);
		kd_tgt.nearestKSearch(*cldptr_src_fpfh, i, n, vec_idx_tgt, vec_dist_tgt);
		for (size_t j = 0; j < vec_idx_tgt.size(); j++) {
			bool b_remove = true;
			int idx = vec_idx_tgt[j];
			std::vector<int> vec_idx_src(n);
			std::vector<float> vec_dist_src(n);
			kd_src.nearestKSearch(*cldptr_tgt_fpfh, idx, n, vec_idx_src, vec_dist_src);
			for (size_t k = 0; k < n; k++) {
				if (vec_idx_src[k] == i) {
					b_remove = false;
					break;
				}
			}
			if (b_remove == false) {
				pcl::Correspondence corr;
				corr.index_query = i;
				corr.index_match = vec_idx_tgt[j];
				corr.distance = vec_dist_tgt[j];
				corrptr->push_back(corr);
				vec_corr_src[i]++;
				vec_corr_tgt[vec_idx_tgt[j]]++;
			}
		}
	}

	t2 = clock();
	//std::cout << std::endl;
	std::cout << std::fixed << std::setprecision(3) << "----correspondenceSearch (s): " << double(t2 - t1) / CLOCKS_PER_SEC << std::endl;
}

void Reg::grorReg(	pcl::PointCloud<PointType>::Ptr& cldptr_src, 
					pcl::PointCloud<PointType>::Ptr& cldptr_tgt, 
					pcl::CorrespondencesPtr& corrptr, 
					const double d_resolution, 
					pcl::PointCloud<PointType>::Ptr& cldptr_gror, 
					Eigen::Matrix4f& transform_gror, 
					double& score_gror) {
	clock_t t1, t2;
	t1 = clock();

	int n_count;
	pcl::registration::GRORInitialAlignment<PointType, PointType, float> gror;
	gror.setInputSource(cldptr_src);
	gror.setInputTarget(cldptr_tgt);
	gror.setResolution(d_resolution);
	gror.setOptimalSelectionNumber(800);  // 800
	//gror.setNumberOfThreads(1);
	gror.setInputCorrespondences(corrptr);
	gror.align(*cldptr_gror);
	n_count = gror.getBestCount();
	score_gror = gror.getFitnessScore();
	transform_gror = gror.getFinalTransformation();

	t2 = clock();
	//std::cout << std::endl;
	std::cout << std::fixed << std::setprecision(3) << "gror score: " << score_gror << ", bestcount: " << n_count << std::endl;
	std::cout << std::fixed << std::setprecision(3) << "----compute gror (s): " << double(t2 - t1) / CLOCKS_PER_SEC << std::endl;
}

void Reg::gicpReg(	pcl::PointCloud<PointType>::Ptr& cldptr_src, 
					pcl::PointCloud<PointType>::Ptr& cldptr_tgt, 
					Eigen::Matrix4f& transform_init, 
					pcl::PointCloud<PointType>::Ptr& cldptr_gicp, 
					Eigen::Matrix4f& transform_gicp, 
					double& score_gicp) {
	clock_t t1, t2;
	t1 = clock();

	pcl::GeneralizedIterativeClosestPoint<PointType, PointType> gicp;
	gicp.setInputSource(cldptr_src);
	gicp.setInputTarget(cldptr_tgt);
	gicp.setMaxCorrespondenceDistance(5.0);  // 5.0
	gicp.setMaximumIterations(100);  // 200
	gicp.setTransformationEpsilon(1e-10);  // 5e-4
	gicp.setEuclideanFitnessEpsilon(0.001);  // -std::numeric_limits<double>::max()
	gicp.align(*cldptr_gicp, transform_init);
	score_gicp = gicp.getFitnessScore();
	transform_gicp = gicp.getFinalTransformation();

	t2 = clock();
	//std::cout << std::endl;
	std::cout << std::fixed << std::setprecision(3) << "gicp has converged: " << gicp.hasConverged() << ", score: " << score_gicp << std::endl;
	std::cout << std::fixed << std::setprecision(3) << "----compute gicp (s): " << double(t2 - t1) / CLOCKS_PER_SEC << std::endl;
}

}  // namespace cldreg

